# CUDA_VISIBLE_DEVICES=2,3,4,5,6,7,8,9 python infermulti.py --cfg dino_cr.py --pth /home/challenge/dataset/mmdet_bisai/work_dirs/dino_cr/epoch_1.pth
import os
import argparse
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
import cv2
from tqdm import tqdm
from mmengine import ConfigDict
from mmengine.config import Config
from mmdet.utils import register_all_modules
from mmdet.apis import init_detector, inference_detector
from ensemble_boxes import weighted_boxes_fusion
import pandas as pd
import time

# 设置只使用 1-8 号显卡
os.environ["CUDA_VISIBLE_DEVICES"] = "2,3,4,5,6,7,8,9"


def get_all_image_paths(root_dir):
    img_paths = []
    for root, _, files in os.walk(root_dir):
        for file in files:
            if file.endswith((".jpg", ".jpeg", ".png")):  # 根据你的图片格式调整
                img_paths.append(os.path.join(root, file))
    return img_paths


def unnorm_box(box, w, h):
    if len(box) == 0:
        return box
    else:
        # 处理一维数组的情况
        box[0] = box[0] * w  # x1 转为像素坐标
        box[1] = box[1] * h  # y1 转为像素坐标
        box[2] = box[2] * w  # x2 转为像素坐标
        box[3] = box[3] * h  # y2 转为像素坐标
        return box


def norm_box(box, w, h):
    if len(box) == 0:
        return box
    else:
        # 处理一维数组
        box[0] = box[0] / w  # x1
        box[1] = box[1] / h  # y1
        box[2] = box[2] / w  # x2
        box[3] = box[3] / h  # y2
        return box


def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)  # 设置当前进程使用的GPU


def cleanup():
    dist.destroy_process_group()


def main_worker(rank, world_size, args):
    setup(rank, world_size)

    # 实际使用的 GPU ID
    device = rank  # 0-7 对应 1-8 号显卡
    model = init_detector(args.cfg, args.pth, device=f"cuda:{device}", cfg_options={})

    # 将模型移动到 GPU 并使用 DDP 包装
    model = model.to(device)
    model = DDP(model, device_ids=[device])

    test_img_dir = "/home/challenge/dataset/testingdata/test_set_A_rename"
    results = []

    for i, file_path in enumerate(tqdm(get_all_image_paths(test_img_dir))):
        img = cv2.imread(file_path)
        h, w = img.shape[:2]

        # DDP 推理
        result = inference_detector(model.module, img)  # 访问原始模型
        pred_instances = result.pred_instances
        boxes = pred_instances.bboxes.cpu().numpy()
        scores = pred_instances.scores.cpu().numpy()
        labels = pred_instances.labels.cpu().numpy()

        # 将框从归一化坐标转为像素坐标
        boxes = [norm_box(box, w, h).tolist() for box in boxes]

        # 加权框融合
        boxes, scores, labels = weighted_boxes_fusion(
            [boxes],
            [scores],
            [labels],
            weights=[1],
            iou_thr=0.25,
            skip_box_thr=0.0001,
        )

        boxes = [unnorm_box(box, w, h) for box in boxes]
        for box, score, label in zip(boxes, scores, labels):
            results.append(
                {
                    "image_path": file_path,
                    "x1": box[0],
                    "y1": box[1],
                    "x2": box[2],
                    "y2": box[3],
                    "score": score,
                    "label": label,
                }
            )
    if len(results) == 0:
        print("No detection results.")
        cleanup()
        return
    if rank == 0:
        # 将结果保存到 CSV 文件
        df = pd.DataFrame(results)
        # filename加上当前时间yyyy/mm/dd/hh/mm
        print("Saving results to CSV file...")
        filename = (
            "./inferres/detection_results"
            + time.strftime("%Y%m%d%H%M", time.localtime())
            + ".csv"
        )
        df.to_csv(filename, index=False)

    cleanup()


def main():
    parser = argparse.ArgumentParser(description="Train a segmentor")
    parser.add_argument("--cfg", type=str, required=True)  # 推理配置文件, 如bisai.py
    parser.add_argument("--pth", type=str, required=True)  # 训好的模型pth
    args = parser.parse_args()

    world_size = 1  # 使用 1-8 号显卡

    # 设置主节点的地址和端口
    os.environ["MASTER_ADDR"] = "localhost"  # 或者使用你的主节点IP地址
    os.environ["MASTER_PORT"] = "12345"  # 确保这个端口没有被占用

    mp.spawn(main_worker, args=(world_size, args), nprocs=world_size, join=True)


if __name__ == "__main__":
    main()
